import os

import numpy as np
from osgeo import gdal, osr
import h5py


def array2tif(output_path, array, geo, nodata):
    cols = array.shape[1]  # obtain cols
    rows = array.shape[0]  # obtain rows

    # 写入影像
    driver = gdal.GetDriverByName('GTiff')
    out_tif = driver.Create(output_path, cols, rows, 1, gdal.GDT_Float32)
    # 设置仿射矩阵
    out_tif.SetGeoTransform(geo)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    out_tif.SetProjection(prj.ExportToWkt())
    # 数据导出
    out_tif.GetRasterBand(1).WriteArray(array)  # 将数据写入内存
    out_tif.GetRasterBand(1).SetNoDataValue(nodata)
    out_tif.FlushCache()  # 将数据写入到硬盘


def main(input_dir, output_dir):
    he5file_list = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith('.he5')]
    rows = 586
    cols = 1383
    for he5file in he5file_list:
        soilMoisture_npd_array = np.ones((rows, cols)) * -9999  # 该文件数据为586*1383，缺失值为-9999
        soilMoisture_sca_array = np.ones((rows, cols)) * -9999  # 该文件数据为586*1383，缺失值为-9999
        latitude_array = np.zeros((rows, cols))
        longitude_array = np.zeros((rows, cols))
        with h5py.File(he5file, "r") as f:
            datafield = f["/HDFEOS/POINTS/AMSR-2 Level 2 Land Data/Data/Combined NPD and SCA Output Fields"]  # 数据存放的地方
            data = datafield[:]  # 数据顺序按照文档，一个元组里面包含35个数据。27197 * 35
            # 每行数据索引第1为纬度，2为经度，3为行数，4为列数，16为SoilMoistureNPD，18为SoilMoistureSCA
            for row in data:
                # print(row[1], row[2], row[3], row[4], row[16], row[18])
                latitude_array[row[3] - 1, row[4] - 1] = row[1]
                longitude_array[row[3] - 1, row[4] - 1] = row[2]
                soilMoisture_npd_array[row[3] - 1, row[4] - 1] = row[16]
                soilMoisture_sca_array[row[3] - 1, row[4] - 1] = row[18]
                # 影像的左上角&右下角坐标
            lonMin, latMax, lonMax, latMin = [longitude_array.min(), latitude_array.max(), longitude_array.max(), latitude_array.min()]
            # 分辨率计算
            latRes = (latMax - latMin) / (rows - 1)
            lonRes = (lonMax - lonMin) / (cols - 1)
            geotransform = (lonMin, lonRes, 0.0, latMax, 0.0, -latRes)
            # 没有子文件夹就创建
            if not os.path.exists(output_dir + "/NPD"):
                os.makedirs(output_dir + "/NPD")
            if not os.path.exists(output_dir + "/SCA"):
                os.makedirs(output_dir + "/SCA")
            output_npd_path = output_dir + "/NPD/" + os.path.splitext(he5file.split("\\")[-1])[0] + '_NPD.tif'
            output_sca_path = output_dir + "/SCA/" + os.path.splitext(he5file.split("\\")[-1])[0] + '_SCA.tif'
            # 写入tif文件
            array2tif(output_npd_path, soilMoisture_npd_array, geotransform, -9999)
            array2tif(output_sca_path, soilMoisture_sca_array, geotransform, -9999)
        f.close()
        print("file:", he5file, "done!")
    print('all done')


if __name__ == '__main__':
    input_dir = r'G:\test\SMAP'
    output_dir = r'G:\test\SMAP\tif'
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)
    main(input_dir, output_dir)
    # test_he5 = r'G:\test\SMAP\AMSR_U2_L2_Land_B02_201207040547_A.he5'

